import labelreg.helpers as helper
from labelreg.helpers import DataReader
import SimpleITK  as sitk
import numpy as np
import sys
import os
from dataprocessor.tools import get_bounding_box
from dataprocessor.itkdatareader import SitkDataReader
import scipy.ndimage
from NIIVisualization.Niiplot import multi_slice_viewer
import nibabel as nib

class NiiDataProcessor():
    def __init__(self,img_dir,lable_dir):
        self.img=SitkDataReader(img_dir)
        self.label=SitkDataReader(lable_dir)
        self.NEW_SIZSE=[96]
        self.LABEL_INDEX=420

    def preprocess(self):

        #get a lable and img
        for i in range(self.img.num_data):
            print("process image %d"%(i+1))
            image=self.img.get_data([i])
            label=self.label.get_data([i])
            print("original size: ",image.shape)
            image, label = self.__crop(image[0, 0:, 0:, 0:, 0], label[0, 0:, 0:, 0:, 0])
            #do the process
            image,label=self.__resize(image,label)
            self.__save(image,label,i)

    def convert(self):
        for i in range(self.label.num_data):
            label = self.label.get_data([i])

    def preprocess_itk(self):
        for i in range(self.img.num_data):
            print("process image %d"%(i))
            if i==13:
                print("bug here")
            image=self.img.get_file_obj(i)
            label=self.label.get_file_obj(i)
            print("original size: ",image.GetSize())

            image,label=self.__sitkresample(image,label)
            # self.__itksave(image, label, i)
            print("image size: ", image.GetSize())
            print("label size: ", label.GetSize())

            image, label = self.__sitkcrop_by_label_index(image, label,self.LABEL_INDEX)

            print("image size: ", image.GetSize())
            print("label size: ", label.GetSize())
            image=self.__itknormalize_itensity(image)
            image,label=self.__itk_padding(image,label)
            print("paddign image size: ", image.GetSize())
            print("paddign label size: ", label.GetSize())
            self.__itksave(image,label,i)

    def __itknormalize_itensity(self,img):
        nor_img=sitk.Normalize(img)

        return nor_img

    def __itksave(self,image,label,i):
        image_path = os.path.join(self.img.dir_name + "_result", os.path.basename(self.img.files[i]))
        label_path = os.path.join(self.label.dir_name + "_result", os.path.basename(self.label.files[i]))
        if not os.path.exists(os.path.dirname(image_path)):
            os.makedirs(os.path.dirname(image_path))
        if not os.path.exists(os.path.dirname(label_path)):
            os.makedirs(os.path.dirname(label_path))
        sitk.WriteImage(image,image_path)
        sitk.WriteImage(label,label_path)

    '''
    对spaceing 重采用，一般大小是1x1x1
    '''
    def __sitkresample(self,image,label,spacing=[1,1,1]):
        newimage = self.__resample_with_interpolator(image,sitk.sitkLinear,spacing)
        newlabel = self.__resample_with_interpolator(label, sitk.sitkNearestNeighbor,spacing)
        return newimage,newlabel

    def __resample_with_interpolator(self, image,interpolator,spacing):
        resample = sitk.ResampleImageFilter()
        resample.SetInterpolator(interpolator)
        resample.SetOutputDirection(image.GetDirection())
        resample.SetOutputOrigin(image.GetOrigin())
        new_spacing = spacing
        resample.SetOutputSpacing(new_spacing)
        orig_size = np.array(image.GetSize(), dtype=np.int)
        orig_spacing = list(image.GetSpacing())
        new_size=[oz*os/nz for oz,os,nz in zip(orig_size,orig_spacing,new_spacing)]
        new_size = np.ceil(new_size).astype(np.int)  # Image dimensions are in integers
        new_size = [int(s) for s in new_size]
        resample.SetSize(new_size)
        newimage = resample.Execute(image)
        return newimage

    def __sitkcrop(self,image,label):



        pass

    def __sitkcrop_by_label_index(self, image, label,label_index):

        array=sitk.GetArrayFromImage(label)
        mask=np.where(array==label_index,1,0)

        # multi_slice_viewer(mask)
        new_label=sitk.GetImageFromArray(mask)
        new_label.CopyInformation(label)#这个函数关键
        center=scipy.ndimage.measurements.center_of_mass(mask)
        print(center)
        pos1,pos2=self.__get_bounding_box(center,mask.shape,size=[self.NEW_SIZSE[0]/2]*3)
        crop_image=image[pos1[2]:pos2[2],pos1[1]:pos2[1],pos1[0]:pos2[0]]
        crop_label=new_label[pos1[2]:pos2[2],pos1[1]:pos2[1],pos1[0]:pos2[0]]

        # crop_image=sitk.GetArrayViewFromImage(image)[pos1[0]:pos2[0],pos1[1]:pos2[1],pos1[2]:pos2[2]]
        # crop_label = mask[pos1[0]:pos2[0], pos1[1]:pos2[1] ,pos1[2]:pos2[2]]
        # multi_slice_viewer(crop_label)
        # return sitk.GetImageFromArray(crop_image),sitk.GetImageFromArray(crop_label)
        return crop_image,crop_label

    def __itk_padding(self,img,lab):
        padder=sitk.ConstantPadImageFilter()
        padder.SetConstant(-3)
        padder.SetPadLowerBound([0,0,0])
        size=self.NEW_SIZSE[0]
        padder.SetPadUpperBound([size-img.GetSize()[0],size-img.GetSize()[1],size-img.GetSize()[2]])
        pad_img=padder.Execute(img)
        padder.SetConstant(0)
        pad_lab=padder.Execute(lab)
        return pad_img,pad_lab

    def __get_bounding_box(self,center,shape,size=[48,48,48]):

        pos1=np.ceil(center)-np.array(size)
        pos2=np.ceil(center)+np.array(size)
        pos1=np.clip(pos1,np.array([0,0,0]),np.array(shape))
        pos2 = np.clip(pos2, np.array([0, 0, 0]), np.array(shape))

        return pos1.astype(int).tolist(),pos2.astype(int).tolist()

    def __crop(self,image,label):

        bbox=get_bounding_box(label)
        crop_img=image[bbox[0].start:bbox[0].stop,bbox[1].start:bbox[1].stop,bbox[2].start:bbox[2].stop]
        crop_label = label[bbox[0].start:bbox[0].stop, bbox[1].start:bbox[1].stop, bbox[2].start:bbox[2].stop]
        # multi_slice_viewer(crop_label)
        # multi_slice_viewer(crop_img)
        print("croped shape: ",crop_label.shape)
        return crop_img,crop_label

    def __resize(self,image,label):

        zoom = [float(nz) / float(oz) for oz, nz in zip(image.shape,self.NEW_SIZSE )]
        new_image = scipy.ndimage.zoom(image, zoom,order=5)
        new_label = scipy.ndimage.zoom(label,zoom,order=0)
        print("scaled size: ",new_label.shape)
        # multi_slice_viewer(new_image)
        # multi_slice_viewer(new_label)
        return  new_image,new_label

    def __convert_label_in_4_dim(self, label):
        intensity = np.sort(np.unique(label))
        print(intensity)
        labels = []
        for i, gray in enumerate(intensity[1:]):
            labels.append(np.copy(label))
            labels[i] = np.where(labels[i] == gray, 1, 0)
        labels_4D = np.stack(labels, -1)
        return  labels_4D

    def __save(self,image,label,i):
        image_path = os.path.join(self.img.dir_name+"_result",os.path.basename(self.img.files[i]))
        label_path = os.path.join(self.label.dir_name+"_result",os.path.basename(self.label.files[i]))
        if not os.path.exists(os.path.dirname(image_path)):
            os.makedirs(os.path.dirname(image_path))
        if not os.path.exists(os.path.dirname(label_path)):
            os.makedirs(os.path.dirname(label_path))
        tmp=nib.Nifti1Image(image,self.img.file_objects[i].affine)
        nib.save(tmp,image_path)
        label_4D = self.__convert_label_in_4_dim(label)
        tmp=nib.Nifti1Image(label_4D,self.label.file_objects[i].affine)
        nib.save(tmp,label_path)



if __name__=="__main__":

    
    pre=NiiDataProcessor("E:\MIA_CODE_DATA\zhuang_data\MMWHS\MRI\mr-image", "E:\MIA_CODE_DATA\zhuang_data\MMWHS\MRI\mr-label")
    pre.preprocess_itk()
    pre=NiiDataProcessor("E:\MIA_CODE_DATA\zhuang_data\MMWHS\CT\\train\ct-image", "E:\MIA_CODE_DATA\zhuang_data\MMWHS\CT\\train\ct-label")
    pre.preprocess_itk()
    # mr_label_reader=DataReader("E:\MIA_CODE_DATA\zhuang_data\MMWHS\MRI\mr-label")
    # mr_img_reader=DataReader("E:\MIA_CODE_DATA\zhuang_data\MMWHS\MRI\mr-image")
    # data=mr_label_reader.get_data(case_indices=[1],label_indices=[1])
    # pre=NiiDataPreprocessor("E:\MIA_CODE_DATA\zhuang_data\MMWHS\MRI\mr-image", \
    #                                  "E:\MIA_CODE_DATA\zhuang_data\MMWHS\MRI\mr-label")

    # data=img_reader.get_data(case_indices=[1])
    # crop_3d_image()
